package com.tuniu.agents.config;

import com.tuniu.agents.advisor.*;
import com.tuniu.agents.agent.AgentFactory;
import com.tuniu.agents.agent.AgentManager;
import com.tuniu.agents.agent.AgentManagerGroup;
import com.tuniu.agents.command.CommandStorage;
import com.tuniu.agents.common.IdGenerator;
import com.tuniu.agents.common.command.RedisCommandStateStorage;
import com.tuniu.agents.common.service.TraceService;
import com.tuniu.agents.common.service.itinerary.DefaultItineraryServiceImpl;
import com.tuniu.agents.common.service.itinerary.ItineraryService;
import com.tuniu.agents.common.util.NumberingUtil;
import com.tuniu.agents.memory.MemoryManager;
import com.tuniu.agents.message.MessageRender;
import com.tuniu.agents.message.TraceLogger;
import com.tuniu.agents.prompt.PromptManager;
import io.netty.channel.ChannelOption;
import jakarta.annotation.PostConstruct;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.http.client.ReactorClientHttpRequestFactory;
import org.springframework.web.client.RestClient;
import reactor.netty.http.client.HttpClient;

import java.util.List;
import java.util.Map;

/**
 * Defines the application configuration
 */
@Configuration
public class ApplicationConfig {

    @Autowired
    @Qualifier("messageRedisTemplate")
    private RedisTemplate<String, Message> messageRedisTemplate;

    @Autowired
    @Qualifier("mapRedisTemplate")
    private RedisTemplate<String, Map<String, Object>> mapRedisTemplate;

    @Autowired
    @Qualifier("stringRedisTemplate")
    private StringRedisTemplate stringRedisTemplate;

    @PostConstruct
    public void init() {
        NumberingUtil.setRedisTemplate(stringRedisTemplate);
    }

    @Bean
    public ChatMemory chatMemory() {
        return new RedisChatMemory(messageRedisTemplate);
    }

    @Bean
    public CommandStorage commandStorage() {
        return new RedisCommandStateStorage(mapRedisTemplate);
    }

    @Bean
    public ChatEnvironmentMemory chatExtendMemory() {
        return new RedisEnvironmentMemory(mapRedisTemplate);
    }

    @Bean
    public AgentManagerGroup agentManagerGroup(ChatClient.Builder modelBuilder, ChatMemory chatMemory, PromptManager promptManager, ChatEnvironmentMemory chatEnvironmentMemory,
                                               AgentManager agentManager, MemoryManager memoryManager, MessageRender messageRender, TraceLogger traceLogger, IdGenerator idGenerator) {
        return new AgentManagerGroup(modelBuilder, chatMemory, promptManager, chatEnvironmentMemory, List.of(new AgentChatMemoryAdvisor(chatMemory),
                new SimpleLoggerAdvisor(chatEnvironmentMemory, traceLogger, idGenerator),
                new ChatExtendMemoryAdvisor(chatEnvironmentMemory)), agentManager, memoryManager, messageRender);
    }

    @Bean
    public AgentFactory agentFactory(AgentManagerGroup agentManagerGroup) {
        return new AgentFactory(agentManagerGroup);
    }

    @Bean
    @ConditionalOnMissingBean
    public RestClient.Builder restClientBuilder() {
        // Create Reactor Netty HttpClient and set timeout
        HttpClient httpClient = HttpClient.create()
                .responseTimeout(java.time.Duration.ofSeconds(30)) // Set response timeout to 30 seconds
                .option(ChannelOption.CONNECT_TIMEOUT_MILLIS, 30000); // Set connection timeout to 30 seconds

        // Create ReactorClientHttpRequestFactory
        ReactorClientHttpRequestFactory requestFactory = new ReactorClientHttpRequestFactory(httpClient);

        // Configure RestClient.Builder
        RestClient.Builder restClientBuilder = RestClient.builder();
        restClientBuilder.requestFactory(requestFactory);

        return restClientBuilder;
    }

    @Bean
    public TraceService traceService() {
        return new TraceService.DefaultTraceService();
    }

    @Bean
    public ItineraryService getItineraryService() {
        return new DefaultItineraryServiceImpl();
    }

}